import pandas as pd
import os


class Dataloader:
    def __init__(self,path):
        self.path = path
        self.dict_code_df = {}

    def get_code_df(self,code):
        if code in self.dict_code_df.keys():
            return self.dict_code_df[code]
        else:
            df = self._load_from_file(code)
            self.dict_code_df[code] = df
        return df

    def get_codes(self, codes):
        dfs = []
        for code in codes:
            df = self.get_code_df(code)
            df.dropna(inplace=True)
            dfs.append(df)
        #calendar = self._calc_calendar(dfs)
        df_all = pd.concat(dfs, axis=0)
        df_all.sort_index(inplace=True)
        return df_all
        #print(df_all[['pct_chg','rate','close','code']])

    def load(self, codes, start_time='20100101', end_time='20211231', fields=[], names=[]):
        from logic.global_objs import E

        if len(fields) > 0:
            for code in codes:
                for feature, name in zip(fields, names):
                    if feature == "":
                        continue
                    expr = E.get_expression(feature)
                    se = expr.load(code)
                    se.name = name
                    df = self.get_code_df(code)
                    df[name] = se
        df_all = self.get_codes(codes)
        df_all = df_all.loc[start_time:end_time]
        return df_all

    def _clean_df(self, df):
        df.rename(columns={'trade_date': 'date', 'ts_code':'code','vol':'volume'}, inplace=True)
        df['date'] = df['date'].apply(lambda x:str(x))
        df = df[['open','high','low','close','date','code','volume']]
        df.index = df['date']
        df.sort_index(ascending=True,inplace=True)
        df['rate'] = df['close'].pct_change()
        return df

    def _load_from_file(self,code):
        filename = '{}/{}.csv'.format(self.path, code)
        if os.path.exists(filename):
            df = pd.read_csv(filename, index_col=[0]) #index_col指定是第一列，这是保存的时候生成的。
            df = self._clean_df(df)
        else:
            print('load_from_file error')
            return None
        return df

if __name__ == '__main__':
    from logic.global_objs import D
    from engine.data.data_handler import DataHandler
    #fields, names = DataHandler().get_kbar_fields_names()

    rule = '$close > Max(Ref($high,1),20)'
    rule2= '$close < Min(Ref($low,1),10)'
    fields = [rule,rule2,'Return($close,20)']
    names = fields

    #fields.append('Ref($close,-1)/$close -1')
    #names.append('label')

    #fields.append('QCut(label,10)')
    #names.append('label_c')


    df_all = D.load(['000300.SH'], start_time='20100101', fields=fields, names=names)

    #se = df_all[[rule,rule2]].sum(axis=1)
    #df_all['to_buy'] = se > 1

    print(df_all.head(30))
    print(df_all.tail(30))
    print(df_all[['Return($close,20)']])
